import mxnet as mx
import mxnet.ndarray as nd
import mxnet.gluon.nn as nn
import mxnet.autograd as ag
from .params import load_param
import numpy as np
import logging
from models.defomableconvnets.config import config
class SymbolBlock(nn.Block):
    def __init__(self, sym, input_names= ("data",), resume = None, pretrained = None):
        super(SymbolBlock,self).__init__()
        self.features = mx.gluon.nn.SymbolBlock(
            sym,
            [mx.sym.Variable(x) for x in input_names]
        )

        net_params = self.features.collect_params()
        if resume is not None:
            args, auxes = load_param(resume)
            for param in args.keys():
                if param in net_params:
                    net_params[param]._load_init(args[param], ctx=mx.cpu())
            for param in auxes.keys():
                if param in net_params:
                    net_params[param]._load_init(auxes[param], ctx=mx.cpu())

        if pretrained is not None:

            means = np.tile(np.array(config.TRAIN.BBOX_MEANS),
                            2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)
            stds = np.tile(np.array(config.TRAIN.BBOX_STDS), 2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)
            args_params = nd.load("output/rfcn_voc-9-0.814870423711.params")
            weight = args_params['rfcn_bbox_weight']
            bias = args_params['rfcn_bbox_bias']
            repeat = bias.shape[0] / means.shape[0]

            args_params['rfcn_bbox_weight'] = weight * mx.nd.repeat(mx.nd.array(stds), repeats=repeat).reshape(
                (bias.shape[0], 1, 1, 1))
            args_params['rfcn_bbox_bias'] = args_params['rfcn_bbox_bias'] * mx.nd.repeat(mx.nd.array(stds),
                                                                                         repeats=repeat) + mx.nd.repeat(
                mx.nd.array(means), repeats=repeat)

            for param in args_params.keys():
                net_params = self.features.collect_params()
                if param in net_params:
                    net_params[param]._load_init(args_params[param], ctx=mx.cpu())

    def forward(self, *args):
        return self.features(*args)


class FPNROIPooling(nn.HybridBlock):
    def __init__(self,feat_strides=(4,8,16,32),
                 pooled_height=7, pooled_width=7, with_deformable=True, output_dim=256):
        super(FPNROIPooling,self).__init__()
        self.pooled_height = pooled_height
        self.pooled_width = pooled_width
        self.feat_strides = feat_strides
        self.with_deformable = with_deformable
        self.output_dim = output_dim
        self.in_grad_hist_list = []
        self.num_strides = len(self.feat_strides)

        self.fc1 = nn.Dense(7*7*2)
        self.fc2 = nn.Dense(7*7*2)
        self.fc3 = nn.Dense(7*7*2)
        self.fc4 = nn.Dense(7*7*2)

        self.fc1.collect_params().initialize(init = mx.init.Zero())
        self.fc2.collect_params().initialize(init = mx.init.Zero())
        self.fc3.collect_params().initialize(init = mx.init.Zero())
        self.fc4.collect_params().initialize(init = mx.init.Zero())

        self.fpn_fc = [self.fc1,self.fc2,self.fc3,self.fc4]
        self.fc1.collect_params().setattr('lr_mult', 0.01)
        self.fc2.collect_params().setattr('lr_mult', 0.01)
        self.fc3.collect_params().setattr('lr_mult', 0.01)
        self.fc4.collect_params().setattr('lr_mult', 0.01)

    def hybrid_forward(self, F, rois_nd, fpn_p2, fpn_p3,fpn_p4,fpn_p5):
        _, x0, y0, x1, y1 = F.split(data=rois_nd, num_outputs=5, axis=1, squeeze_axis=1)
        w = x1 - x0 + 1
        h = y1 - y0 + 1
        feat_id = F.clip(F.floor(2 + F.log2(F.sqrt(w * h) / 224)), 0, len(self.feat_strides) - 1)
        in_data = [fpn_p2, fpn_p3, fpn_p4, fpn_p5]
        rois_p = [rois_nd,rois_nd,rois_nd,rois_nd]
        roi_pool = [None for _ in range(self.num_strides)]
        for i in range(self.num_strides):
            roi_offset_t = F.contrib.DeformablePSROIPooling(data=in_data[i],
                                                                rois=rois_p[i],
                                                                group_size=1, pooled_size=7,
                                                                sample_per_part=4, no_trans=True,
                                                                part_size=7, output_dim=256,
                                                                spatial_scale=1.0 / self.feat_strides[i])
            roi_offset = self.fpn_fc[i](roi_offset_t)
            roi_offset_reshape = F.reshape(data=roi_offset, shape=(-1, 2, 7, 7))
            roi_pool[i] = F.contrib.DeformablePSROIPooling(data=in_data[i],
                                                                rois=rois_p[i],
                                                                trans=roi_offset_reshape,
                                                                group_size=1, pooled_size=7,
                                                                sample_per_part=4, no_trans=False,
                                                                part_size=7,
                                                                output_dim=self.output_dim,
                                                                spatial_scale=1.0 / self.feat_strides[
                                                                    i], trans_std=0.1)
        r = F.where(feat_id == 0, roi_pool[0], roi_pool[1])
        r = F.where(feat_id == 1, roi_pool[1], r)
        r = F.where(feat_id == 2, roi_pool[2], r)
        r = F.where(feat_id == 3, roi_pool[3], r)
        return r


class FasterRCNN(nn.Block):
    def __init__(self, cfg, rpn_rcnn_feat, *args, **kwargs):
        super(FasterRCNN, self).__init__(*args, **kwargs)
        self.cfg = cfg
        with self.name_scope():
            # Backbone
            self.rpn_rcnn_feat = rpn_rcnn_feat

            # RPN
            self.rpn_relu = nn.Conv2D(channels=512, kernel_size=(3, 3), strides=(1, 1), padding=(1, 1), activation="relu", prefix="rpn_conv_3x3_")
            self.rpn_cls_score = nn.Conv2D(channels=2 * cfg.network.NUM_ANCHORS, kernel_size=(1, 1), strides=(1, 1), prefix="rpn_cls_score_")
            self.rpn_bbox_pred = nn.Conv2D(channels=4 * cfg.network.NUM_ANCHORS, kernel_size=(1, 1), strides=(1, 1), prefix="rpn_bbox_pred_")

            # RCNN
            self.relu_new_1 = nn.Conv2D(channels=1024, kernel_size=(1, 1), strides=(1, 1), activation="relu", prefix="conv_new_1_")
            self.relu_new_1.collect_params().setattr('lr_mult', 3)
            self.rfcn_cls = nn.Conv2D(channels=7 * 7 * cfg.dataset.NUM_CLASSES, kernel_size=(1, 1), strides=(1, 1), prefix="rfcn_cls_")
            self.rfcn_bbox = nn.Conv2D(channels=7 * 7 * (2 if cfg.CLASS_AGNOSTIC else cfg.dataset.NUM_CLASSES) * 4, kernel_size=(1, 1), strides=(1, 1), prefix="rfcn_bbox_")
            self.rfcn_cls_offset_t = nn.Conv2D(channels=2 * 7 * 7 * cfg.dataset.NUM_CLASSES, kernel_size=(1, 1), strides=(1, 1), prefix="rfcn_cls_offset_t_")
            self.rfcn_bbox_offset_t = nn.Conv2D(channels=2 * 7 * 7, kernel_size=(1, 1), strides=(1, 1), prefix="rfcn_bbox_offset_t_")

    def forward(self, data, im_info, gt_boxes = None):
        cfg = self.cfg
        num_classes = cfg.dataset.NUM_CLASSES
        num_reg_classes = (2 if cfg.CLASS_AGNOSTIC else num_classes)
        num_anchors = cfg.network.NUM_ANCHORS

        # RPN
        rpn_feat, rcnn_feat = self.rpn_rcnn_feat(data)
        rpn_relu = self.rpn_relu(rpn_feat)
        rpn_cls_score = self.rpn_cls_score(rpn_relu)
        rpn_bbox_pred = self.rpn_bbox_pred(rpn_relu)
        rois = self.proposal(rpn_cls_score, rpn_bbox_pred, gt_boxes, im_info)

        # RCNN
        relu_new_1 = self.relu_new_1(rcnn_feat)
        rfcn_cls = self.rfcn_cls(relu_new_1)
        rfcn_bbox = self.rfcn_bbox(relu_new_1)

        # Deformable PSROIPooling
        rfcn_cls_offset_t = self.rfcn_cls_offset_t(relu_new_1)
        rfcn_bbox_offset_t = self.rfcn_bbox_offset_t(relu_new_1)
        rfcn_cls_offset = mx.nd.contrib.DeformablePSROIPooling(data=rfcn_cls_offset_t, rois=rois, group_size=7, pooled_size=7, sample_per_part=4, no_trans=True, part_size=7, output_dim=2 * num_classes, spatial_scale=1.0 / self.cfg.network.RPN_FEAT_STRIDE)
        rfcn_bbox_offset = mx.nd.contrib.DeformablePSROIPooling(data=rfcn_bbox_offset_t, rois=rois, group_size=7, pooled_size=7, sample_per_part=4, no_trans=True, part_size=7, output_dim=2, spatial_scale=1.0 / self.cfg.network.RPN_FEAT_STRIDE)
        psroipooled_cls_rois = mx.nd.contrib.DeformablePSROIPooling(data=rfcn_cls, rois=rois, trans=rfcn_cls_offset, group_size=7, pooled_size=7, sample_per_part=4, no_trans=False, trans_std=0.1, output_dim=num_classes, spatial_scale=1.0 / self.cfg.network.RPN_FEAT_STRIDE, part_size=7)
        psroipooled_loc_rois = mx.nd.contrib.DeformablePSROIPooling(data=rfcn_bbox, rois=rois, trans=rfcn_bbox_offset, group_size=7, pooled_size=7, sample_per_part=4, no_trans=False, trans_std=0.1, output_dim=8, spatial_scale=1.0 / self.cfg.network.RPN_FEAT_STRIDE, part_size=7)
        rcnn_cls_score = mx.nd.Pooling(data=psroipooled_cls_rois, pool_type='avg', global_pool=True, kernel=(7, 7))
        rcnn_bbox_pred = mx.nd.Pooling(data=psroipooled_loc_rois, pool_type='avg', global_pool=True, kernel=(7, 7))
        return rois, rpn_cls_score, rpn_bbox_pred, rcnn_cls_score, rcnn_bbox_pred

    def proposal(self, rpn_cls_score, rpn_bbox_pred, gt_boxes, im_info):
        # im_info = mx.nd.array([[data.shape[2], data.shape[3], 1]], ctx = rpn_cls_score.context)
        cfg = self.cfg
        rpn_cls_score = mx.nd.reshape(rpn_cls_score, shape=(0, 2, -1, 0)).softmax(axis = 1).reshape(shape = rpn_cls_score.shape)
        rois = mx.nd.contrib.Proposal(
            cls_prob=rpn_cls_score, bbox_pred=rpn_bbox_pred, im_info=im_info, name='rois',
            feature_stride=cfg.network.RPN_FEAT_STRIDE, scales=tuple(cfg.network.ANCHOR_SCALES),
            ratios=tuple(cfg.network.ANCHOR_RATIOS),
            rpn_pre_nms_top_n=cfg.TRAIN.RPN_PRE_NMS_TOP_N, rpn_post_nms_top_n=cfg.TRAIN.RPN_POST_NMS_TOP_N,
            threshold=cfg.TRAIN.RPN_NMS_THRESH, rpn_min_size=cfg.TRAIN.RPN_MIN_SIZE)
        if gt_boxes is not None:
            gt_rois_nb = nd.zeros(shape=(gt_boxes.shape[1], 1))
            rois = nd.concat(gt_rois_nb, gt_boxes[0, :,:4], dim = 1)
        return rois
